#encoding=utf-8
'''
Created on 2011-11-14

@author: guanyw

python version:2.6
'''

from corpus.loader import LazyCorpusLoaderEx
from corpus.writer import write_pairs
from corpus.pair_reader import RTECorpusReaderEx, norm




    
class RTEPairSimple:
    '''RTE Pair: Simplify RTEPair'''
    
    def __init__(self, pair,doc_id=None,topic_id=None,text_id=None,
                 hyp_id=None,text=None,hyp=None,value=None):
        """
        @param doc_id: the text document id
        @param topic_id : topic id 
        @param text: the text component of the pair
        @param hyp: the hypothesis component of the pair
        @param value: classification label for the pair
        """
        def set_field(field_name,field_value):
            if field_name in pair.attrib:
                return pair.attrib[field_name]
            return field_value
        
        self.topic_id = set_field('t_id', topic_id)
        self.doc_id = set_field('doc_id', doc_id)
        self.text_id = set_field('s_id', text_id)
        self.hyp_id = set_field('h_id', hyp_id)
        self.value = set_field('value', value)
        self.value = set_field('tag', self.value)
        if self.value:
            self.value = norm(self.value)
        if self.text_id and not text:
            self.text = pair.text
        else:
            self.text = text
            
        if self.hyp_id and not hyp:
            self.hyp = pair.text.strip()
        else:
            self.hyp = hyp
    
    
    def __repr__(self):
        return '\tt_id:%s; doc_id:%s; text_id:%s; hyp_id:%s;value:%s;text:%s;hyp:%s\n' \
    %(self.topic_id,self.doc_id,self.text_id,self.hyp_id,self.value,self.text,self.hyp)
        
 

class RTECorpusReaderEx2(RTECorpusReaderEx):
    """
    Corpus reader for RTE corpus.
    
    
    TWO FORMAT: 
    hypothesis :
    <TOPIC t_id="1">
        <H h_id="11">...</H>
        <H h_id="12">...</H>
    </TOPIC>
    <TOPIC t_id="2">
        <H h_id="21">...</H>
    </TOPIC>
    
    text :
    <document>
        <s s_id="0">...</s>
        <s s_id="0">...</s>
    </document>
    """
        
    def _read_etree(self, doc):
        """
        Map the XML input into an RTEPair.
        
        This uses the C{getiterator()} method from the ElementTree package to
        find all the C{<pair>} elements.
        
        @param doc: a parsed XML document
        @rtype: C{list} of L{RTEPair}s
        """
        for pair in doc.getiterator('s'):
            if 's_id' in pair.attrib:

                return [RTEPairSimple(pair) for pair in doc.getiterator("s")]
            
        ret_list = []
        for pair in doc.getiterator('TOPIC'):
            topic_id = pair.attrib['t_id']
            for hyp_pair in pair:

                for doc_id in range(1,11):
                    text_pair = self.pairs('topics/%s/%s.xml'%(topic_id,doc_id))
                                        
                    for tp in text_pair:
                        ret_list.append(RTEPairSimple(hyp_pair,
                                                      topic_id=topic_id,
                                                      doc_id=doc_id,
                                                      text=tp.text,
                                                      text_id=tp.text_id))
            
        return ret_list


class RTECorpusReaderEx3(RTECorpusReaderEx):
    """
    Corpus reader for RTE corpus.
    
    
    FORMAT: 
    <entailment_corpus>
      <TOPIC t_id="1">
        <H h_id="11">
          <H_sentence>hypothesis</H_sentence>
          <text s_id="0" doc_id="1" value="N">text</text>
          <text s_id="2" doc_id="1" value="N">text</text>
            ...
        </H>
      </TOPIC>
      <TOPIC t_id="2">
        <H h_id="21">
          <H_sentence>hypothesis</H_sentence>
          <text s_id="0" doc_id="1" value="N">text</text>
            ...
        </H>
      </TOPIC>
    </entailment_corpus>
    """
        
    def _read_etree(self, doc):
        """
        Map the XML input into an RTEPair.
        
        This uses the C{getiterator()} method from the ElementTree package to
        find all the C{<pair>} elements.
        
        @param doc: a parsed XML document
        @rtype: C{list} of L{RTEPairSimple}s
        """
            
        ret_list = []
        for pair in doc.getiterator('TOPIC'):
            topic_id = pair.attrib['t_id']
            for hyp_pair in pair:
                hyp_id = hyp_pair.attrib['h_id']
                hyp = hyp_pair[0].text

                if len(hyp_pair) >1:
                    for pair in hyp_pair[1:]:
                        ret_list.append(RTEPairSimple(pair,topic_id=topic_id,hyp_id=hyp_id,hyp=hyp))
        
        return ret_list

rte_train_nto1_labeled = LazyCorpusLoaderEx('../../corpus/data/train', RTECorpusReaderEx3, r'(?!\.).*\.xml')

if __name__ == '__main__':
    rte = LazyCorpusLoaderEx('train', RTECorpusReaderEx2, r'(?!\.).*\.xml')
    #topics/1/1.xml
    pairs1 = rte.pairs(['rte_hypothesis.xml'])
#    print pairs1
    write_pairs(pairs1,'1.xml',verbose='full')
    
    pairs2 = LazyCorpusLoaderEx('train', RTECorpusReaderEx3, r'(?!\.).*\.xml').pairs(['rte.xml'])
#    print pairs2[1:20]
    
    for i in range(len(pairs1)):
        if str(pairs1[i]) != str(pairs2[i]):
            print pairs1[1]
    write_pairs(pairs2,'2.xml',verbose='full')
